"""
Twiss参数相位空间图
展示alpha > 0, alpha = 0, alpha < 0三种情况下的粒子分布椭圆
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse

# 设置中文字体（用于print）
plt.rcParams['axes.unicode_minus'] = False

# 固定参数
epsilon_rms = 1.0  # RMS emittance (归一化)

# 创建1x3子图
fig, axes = plt.subplots(1, 3, figsize=(18, 6))

# 定义三种情况
cases = [
    {'alpha': 0.8, 'label': r'$\alpha > 0$', 'title': '(a) Converging Beam', 
     'description': 'Beam converging toward focus'},
    {'alpha': 0.0, 'label': r'$\alpha = 0$', 'title': '(b) Minimum Spot Size', 
     'description': 'Beam at minimum spot size'},
    {'alpha': -0.8, 'label': r'$\alpha < 0$', 'title': '(c) Diverging Beam', 
     'description': 'Beam diverging from focus'}
]

for idx, case in enumerate(cases):
    ax = axes[idx]
    alpha = case['alpha']
    
    # 根据约束条件 beta*gamma - alpha^2 = 1 选择beta和gamma
    # 选择beta使得椭圆形状合理
    if alpha > 0:
        beta = 1.5  # 收敛时beta较大
    elif alpha == 0:
        beta = 1.0  # 最小斑点时beta适中
    else:
        beta = 0.8  # 发散时beta较小
    
    # 计算gamma: gamma = (1 + alpha^2) / beta
    gamma = (1 + alpha**2) / beta
    
    # 验证约束: beta*gamma - alpha^2 = 1
    assert abs(beta * gamma - alpha**2 - 1) < 1e-10, "Constraint violated!"
    
    # 计算椭圆参数
    # 椭圆方程: gamma*x^2 + 2*alpha*x*x' + beta*x'^2 = epsilon_rms
    # 转换为标准椭圆形式需要特征值分解
    
    # 构建协方差矩阵
    sigma_matrix = epsilon_rms * np.array([[beta, -alpha], 
                                          [-alpha, gamma]])
    
    # 特征值分解
    eigenvals, eigenvecs = np.linalg.eigh(sigma_matrix)
    angle = np.degrees(np.arctan2(eigenvecs[1, 0], eigenvecs[0, 0]))
    
    # 椭圆半轴长度
    a = np.sqrt(eigenvals[0])  # 较长半轴
    b = np.sqrt(eigenvals[1])  # 较短半轴
    
    # 绘制椭圆
    ellipse = Ellipse((0, 0), 2*a, 2*b, angle=angle, 
                     fill=False, edgecolor='blue', linewidth=2.5, linestyle='-')
    ax.add_patch(ellipse)
    
    # 绘制粒子点（模拟粒子分布）
    n_particles = 200
    # 生成符合椭圆分布的粒子
    theta = np.random.uniform(0, 2*np.pi, n_particles)
    r = np.random.uniform(0, 1, n_particles)
    # 椭圆内的均匀分布
    x_sample = np.sqrt(r) * np.cos(theta) * a
    x_prime_sample = np.sqrt(r) * np.sin(theta) * b
    
    # 旋转到正确的方向
    cos_angle = np.cos(np.radians(angle))
    sin_angle = np.sin(np.radians(angle))
    x_rot = x_sample * cos_angle - x_prime_sample * sin_angle
    x_prime_rot = x_sample * sin_angle + x_prime_sample * cos_angle
    
    # 添加一些随机噪声使其更真实
    noise_factor = 0.1
    x_rot += np.random.normal(0, noise_factor * a, n_particles)
    x_prime_rot += np.random.normal(0, noise_factor * b, n_particles)
    
    ax.scatter(x_rot, x_prime_rot, s=10, alpha=0.6, color='red', edgecolors='darkred', linewidths=0.5)
    
    # 绘制坐标轴
    ax.axhline(y=0, color='k', linestyle='--', linewidth=0.8, alpha=0.5)
    ax.axvline(x=0, color='k', linestyle='--', linewidth=0.8, alpha=0.5)
    
    # 标注参数
    param_text = f'α = {alpha:.2f}\nβ = {beta:.2f}\nγ = {gamma:.2f}\nε = {epsilon_rms:.2f}'
    ax.text(0.02, 0.98, param_text, transform=ax.transAxes, 
           fontsize=11, verticalalignment='top',
           bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))
    
    # 标注椭圆方程
    eq_text = r'$\gamma x^2 + 2\alpha xx^\prime + \beta x^{\prime 2} = \epsilon$'
    ax.text(0.5, 0.02, eq_text, transform=ax.transAxes, 
           fontsize=11, ha='center',
           bbox=dict(boxstyle='round', facecolor='lightblue', alpha=0.7))
    
    # 设置标题和标签
    ax.set_title(case['title'], fontsize=14, fontweight='bold')
    ax.set_xlabel(r'$x$ [mm]', fontsize=12, fontweight='bold')
    ax.set_ylabel(r"$x'$ [mm·rad]", fontsize=12, fontweight='bold')
    
    # 设置坐标轴范围（根据椭圆大小自动调整）
    margin = 0.2
    x_range = max(np.abs(x_rot)) * (1 + margin)
    x_prime_range = max(np.abs(x_prime_rot)) * (1 + margin)
    ax.set_xlim([-x_range, x_range])
    ax.set_ylim([-x_prime_range, x_prime_range])
    
    ax.grid(True, alpha=0.3)
    ax.set_aspect('equal', adjustable='box')

# 添加总标题
plt.suptitle('Twiss Parameter Phase Space Diagrams', 
             fontsize=16, fontweight='bold', y=1.02)

plt.tight_layout()
plt.savefig('twiss_phase_space.png', dpi=300, bbox_inches='tight')
print("Twiss相位空间图已保存: twiss_phase_space.png")
plt.close()

